{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "8JSGdaDHc_f4"
      },
      "outputs": [],
      "source": [
        "# Copyright 2019 The TensorFlow Hub Authors. All Rights Reserved.\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     http://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License.\n",
        "# =============================================================================="
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "z2_BHI6XdJ30"
      },
      "source": [
        "# Text-to-Video retrieval with S3D MIL-NCE"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Rm0K9ZTgfISB"
      },
      "source": [
        "\u003ctable class=\"tfo-notebook-buttons\" align=\"left\"\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://www.tensorflow.org/hub/tutorials/text_to_video_retrieval_with_s3d_milnce\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" /\u003eView on TensorFlow.org\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/hub/blob/master/examples/colab/text_to_video_retrieval_with_s3d_milnce.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" /\u003eRun in Google Colab\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://github.com/tensorflow/hub/blob/master/examples/colab/text_to_video_retrieval_with_s3d_milnce.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" /\u003eView source on GitHub\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca href=\"https://storage.googleapis.com/tensorflow_docs/hub/examples/colab/text_to_video_retrieval_with_s3d_milnce.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/download_logo_32px.png\" /\u003eDownload notebook\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "\u003c/table\u003e"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "bC_xJPpQd-LO"
      },
      "outputs": [],
      "source": [
        "!pip install -q opencv-python\n",
        "\n",
        "import os\n",
        "\n",
        "import tensorflow.compat.v2 as tf\n",
        "import tensorflow_hub as hub\n",
        "\n",
        "import numpy as np\n",
        "import cv2\n",
        "from IPython import display\n",
        "import math"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "ZxwaK-jf7qkW"
      },
      "source": [
        "## Import TF-Hub model\n",
        "\n",
        "This tutorial demonstrates how to use the [S3D MIL-NCE model](https://tfhub.dev/deepmind/mil-nce/s3d/1) from TensorFlow Hub to do **text-to-video retrieval** to find the most similar videos for a given text query.\n",
        "\n",
        "The model has 2 signatures, one for generating *video embeddings* and one for generating *text embeddings*. We will use these embedding to find the nearest neighbors in the embedding space."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "nwv4ZQ4qmak5"
      },
      "outputs": [],
      "source": [
        "# Load the model once from TF-Hub.\n",
        "hub_handle = 'https://tfhub.dev/deepmind/mil-nce/s3d/1'\n",
        "hub_model = hub.load(hub_handle)\n",
        "\n",
        "def generate_embeddings(model, input_frames, input_words):\n",
        "  \"\"\"Generate embeddings from the model from video frames and input words.\"\"\"\n",
        "  # Input_frames must be normalized in [0, 1] and of the shape Batch x T x H x W x 3\n",
        "  vision_output = model.signatures['video'](tf.constant(tf.cast(input_frames, dtype=tf.float32)))\n",
        "  text_output = model.signatures['text'](tf.constant(input_words))\n",
        "  return vision_output['video_embedding'], text_output['text_embedding']"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "EOZzu9ddekEj"
      },
      "outputs": [],
      "source": [
        "# @title Define video loading and visualization functions  { display-mode: \"form\" }\n",
        "\n",
        "# Utilities to open video files using CV2\n",
        "def crop_center_square(frame):\n",
        "  y, x = frame.shape[0:2]\n",
        "  min_dim = min(y, x)\n",
        "  start_x = (x // 2) - (min_dim // 2)\n",
        "  start_y = (y // 2) - (min_dim // 2)\n",
        "  return frame[start_y:start_y+min_dim,start_x:start_x+min_dim]\n",
        "\n",
        "\n",
        "def load_video(video_url, max_frames=32, resize=(224, 224)):\n",
        "  path = tf.keras.utils.get_file(os.path.basename(video_url)[-128:], video_url)\n",
        "  cap = cv2.VideoCapture(path)\n",
        "  frames = []\n",
        "  try:\n",
        "    while True:\n",
        "      ret, frame = cap.read()\n",
        "      if not ret:\n",
        "        break\n",
        "      frame = crop_center_square(frame)\n",
        "      frame = cv2.resize(frame, resize)\n",
        "      frame = frame[:, :, [2, 1, 0]]\n",
        "      frames.append(frame)\n",
        "\n",
        "      if len(frames) == max_frames:\n",
        "        break\n",
        "  finally:\n",
        "    cap.release()\n",
        "  frames = np.array(frames)\n",
        "  if len(frames) \u003c max_frames:\n",
        "    n_repeat = int(math.ceil(max_frames / float(len(frames))))\n",
        "    frames = frames.repeat(n_repeat, axis=0)\n",
        "  frames = frames[:max_frames]\n",
        "  return frames / 255.0\n",
        "\n",
        "def display_video(urls):\n",
        "    html = '\u003ctable\u003e'\n",
        "    html += '\u003ctr\u003e\u003cth\u003eVideo 1\u003c/th\u003e\u003cth\u003eVideo 2\u003c/th\u003e\u003cth\u003eVideo 3\u003c/th\u003e\u003c/tr\u003e\u003ctr\u003e'\n",
        "    for url in urls:\n",
        "        html += '\u003ctd\u003e'\n",
        "        html += '\u003cimg src=\"{}\" height=\"224\"\u003e'.format(url)\n",
        "        html += '\u003c/td\u003e'\n",
        "    html += '\u003c/tr\u003e\u003c/table\u003e'\n",
        "    return display.HTML(html)\n",
        "\n",
        "def display_query_and_results_video(query, urls, scores):\n",
        "  \"\"\"Display a text query and the top result videos and scores.\"\"\"\n",
        "  sorted_ix = np.argsort(-scores)\n",
        "  html = ''\n",
        "  html += '\u003ch2\u003eInput query: \u003ci\u003e{}\u003c/i\u003e \u003c/h2\u003e\u003cdiv\u003e'.format(query)\n",
        "  html += 'Results: \u003cdiv\u003e'\n",
        "  html += '\u003ctable\u003e'\n",
        "  html += '\u003ctr\u003e\u003cth\u003eRank #1, Score:{:.2f}\u003c/th\u003e'.format(scores[sorted_ix[0]])\n",
        "  html += '\u003cth\u003eRank #2, Score:{:.2f}\u003c/th\u003e'.format(scores[sorted_ix[1]])\n",
        "  html += '\u003cth\u003eRank #3, Score:{:.2f}\u003c/th\u003e\u003c/tr\u003e\u003ctr\u003e'.format(scores[sorted_ix[2]])\n",
        "  for i, idx in enumerate(sorted_ix):\n",
        "    url = urls[sorted_ix[i]];\n",
        "    html += '\u003ctd\u003e'\n",
        "    html += '\u003cimg src=\"{}\" height=\"224\"\u003e'.format(url)\n",
        "    html += '\u003c/td\u003e'\n",
        "  html += '\u003c/tr\u003e\u003c/table\u003e'\n",
        "  return html\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "Ime5V4kDewh8"
      },
      "outputs": [],
      "source": [
        "# @title Load example videos and define text queries  { display-mode: \"form\" }\n",
        "\n",
        "video_1_url = 'https://upload.wikimedia.org/wikipedia/commons/b/b0/YosriAirTerjun.gif' # @param {type:\"string\"}\n",
        "video_2_url = 'https://upload.wikimedia.org/wikipedia/commons/e/e6/Guitar_solo_gif.gif' # @param {type:\"string\"}\n",
        "video_3_url = 'https://upload.wikimedia.org/wikipedia/commons/3/30/2009-08-16-autodrift-by-RalfR-gif-by-wau.gif' # @param {type:\"string\"}\n",
        "\n",
        "video_1 = load_video(video_1_url)\n",
        "video_2 = load_video(video_2_url)\n",
        "video_3 = load_video(video_3_url)\n",
        "all_videos = [video_1, video_2, video_3]\n",
        "\n",
        "query_1_video = 'waterfall' # @param {type:\"string\"}\n",
        "query_2_video = 'playing guitar' # @param {type:\"string\"}\n",
        "query_3_video = 'car drifting' # @param {type:\"string\"}\n",
        "all_queries_video = [query_1_video, query_2_video, query_3_video]\n",
        "all_videos_urls = [video_1_url, video_2_url, video_3_url]\n",
        "display_video(all_videos_urls)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "NCLKv_L_8Anc"
      },
      "source": [
        "## Demonstrate text to video retrieval\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "9oX8ItFUjybi"
      },
      "outputs": [],
      "source": [
        "# Prepare video inputs.\n",
        "videos_np = np.stack(all_videos, axis=0)\n",
        "\n",
        "# Prepare text input.\n",
        "words_np = np.array(all_queries_video)\n",
        "\n",
        "# Generate the video and text embeddings.\n",
        "video_embd, text_embd = generate_embeddings(hub_model, videos_np, words_np)\n",
        "\n",
        "# Scores between video and text is computed by dot products.\n",
        "all_scores = np.dot(text_embd, tf.transpose(video_embd))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "d4AwYmODmE9Y"
      },
      "outputs": [],
      "source": [
        "# Display results.\n",
        "html = ''\n",
        "for i, words in enumerate(words_np):\n",
        "  html += display_query_and_results_video(words, all_videos_urls, all_scores[i, :])\n",
        "  html += '\u003cbr\u003e'\n",
        "display.HTML(html)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "Text-to-Video retrieval with S3D MIL-NCE",
      "private_outputs": true,
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
